from kafka import KafkaProducer
from pyspark.streaming import StreamingContext
#from pyspark.streaming.kafka import KafkaUtils
from pyspark import SparkConf, SparkContext
import json
import sys
from pyspark.sql import DataFrame
from pyspark.sql import SparkSession
from pyspark.sql.functions import window
from pyspark.sql.types import StructType, StructField
from pyspark.sql.types import TimestampType, StringType
from pyspark.sql.functions import col, column, expr
 
def KafkaWordCount(zkQuorum, group, topics, numThreads):
 
    spark = SparkSession \
    .builder \
    .appName("KafkaWordCount") \
    .getOrCreate()
 
    spark.sparkContext.setLogLevel("ERROR")
 
    topicAry = topics.split(",")
        # 将topic转换为hashmap形式，而python中字典就是一种hashmap
    topicMap = {}
    for topic in topicAry:
        topicMap[topic] = numThreads
    #lines = KafkaUtils.createStream(ssc, zkQuorum, group, topicMap).map(lambda x : x[1])
 
    df = spark \
      .readStream \
      .format("kafka") \
      .option("kafka.bootstrap.servers", "localhost:9092") \
      .option("subscribe", "sex") \
      .option("failOnDataLoss","false")\
      .load()
    df.selectExpr( "CAST(timestamp AS timestamp)","CAST(value AS STRING)")
    #lines = df.selectExpr("CAST(value AS STRING)")
 
    windowedCounts = df \
        .withWatermark("timestamp", "1 seconds") \
        .groupBy(
            window(col("timestamp"), "1 seconds" ,"1 seconds"),
            col("value")) \
        .count()
    wind = windowedCounts.selectExpr( "CAST(value AS STRING)","CAST(count AS STRING)")
 
    query = wind.writeStream.option("checkpointLocation", "/check").outputMode("append").foreach(sendmsg).start()
 
    query.awaitTermination()
    query.stop()
 
 
# 格式转化，将格式变为[{1: 3}]
def Get_dic(row):
    res = []
    #for elm in row:
    tmp = {row[0]: row[1]}
    res.append(tmp)
    print(res)
    return json.dumps(res)
 
 
def sendmsg(row):
    print(row)
    if row.count != 0:
        msg = Get_dic(row)
        # 实例化一个KafkaProducer示例，用于向Kafka投递消息
        producer = KafkaProducer(bootstrap_servers='localhost:9092')
        producer.send("result", msg.encode('utf8'))
        # 很重要，不然不会更新
        producer.flush()
 
 
if __name__ == '__main__':
    # 输入的四个参数分别代表着
    # 1.zkQuorum为zookeeper地址
    # 2.group为消费者所在的组
    # 3.topics该消费者所消费的topics
    # 4.numThreads开启消费topic线程的个数
    if (len(sys.argv) < 5):
        print("Usage: KafkaWordCount <zkQuorum> <group> <topics> <numThreads>")
        exit(1)
    zkQuorum = sys.argv[1]
    group = sys.argv[2]
    topics = sys.argv[3]
    numThreads = int(sys.argv[4])
    print(group, topics)
    KafkaWordCount(zkQuorum, group, topics, numThreads)
 
 